import numpy as np
import os
import cv2
import datetime
import time
import re
import shutil

# 在深度学习的训练程序中，锚点0的负样本为锚点1和锚点n

# videoPath = 'D:/Research/2020ContrastiveLearningForSceneLabel/Data/RingroadTestData/video.avi'
videoPath = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round1/round1.avi'
savePath = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round1/allImageReferenceLabel.txt'


ringRoadStartFrame = 0#2897
ringRoadEndFrame = 65857
watiKeyDeley = 0 # 第一帧是暂停的
bakupWaitKeyDeley = 41 # 1 ~ 100， 速度从慢到快
playOrPause = True # 0 play 1 pause，第一帧是暂停的f
maxSpeed = 100 # 其实这个控制的是最慢能有多慢，慢是没有下界的，但是快是有上界的，这个上界取决于对一帧图像的处理时间

idx_and_labels = []

cap = cv2.VideoCapture(videoPath)
frame_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
cap.set(cv2.CAP_PROP_POS_FRAMES, ringRoadStartFrame)
while(1):
    fno = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
    ret, frame = cap.read()
    if not ret:
        break
    # cv2.rectangle(frame, (250,400),(1050,820),(0,0,255),2)
    # frame = frame[400:820,255:1055]
    originFrame = frame.copy()

    cv2.putText(frame, 'fno %d'%fno,(30,30), 2, 0.8, (0,0,255),2)
    cv2.putText(frame,'speed %d (1~100)'%(maxSpeed + 1 - bakupWaitKeyDeley), (30,60),2,0.8,(0,0,255),2)
    if fno == 0:
        cv2.putText(frame,'Must give label of frame 0 !!!', (30,90),2,0.8,(0,0,255),2)
    cv2.imshow('img',frame)
    key = cv2.waitKeyEx(watiKeyDeley)
    if key == 32: # 空格键，播放或暂停
        playOrPause = not playOrPause
        if not playOrPause:
            watiKeyDeley = bakupWaitKeyDeley
        else:
            bakupWaitKeyDeley = watiKeyDeley
            watiKeyDeley = 0
    elif key >= ord('0') and key <= ord('9'):
        idx_and_labels.append([fno, int(chr(key))])
        print([fno, int(chr(key))])
    elif key == 2490368: # 方向键 上， 加速播放
        cap.set(cv2.CAP_PROP_POS_FRAMES, fno)
        if not playOrPause:
            if watiKeyDeley <= 1:
                watiKeyDeley = 1
            else:
                watiKeyDeley = watiKeyDeley - 1
            bakupWaitKeyDeley = watiKeyDeley
        else:
            if bakupWaitKeyDeley <= 1:
                bakupWaitKeyDeley = 1
            else:
                bakupWaitKeyDeley = bakupWaitKeyDeley - 1
    elif key == 2621440: # 方向键 下， 减速播放
        cap.set(cv2.CAP_PROP_POS_FRAMES, fno)
        if not playOrPause:
            if watiKeyDeley < maxSpeed:
                watiKeyDeley = watiKeyDeley + 1
            else:
                watiKeyDeley = maxSpeed
            bakupWaitKeyDeley = watiKeyDeley
        else:
            if bakupWaitKeyDeley < maxSpeed:
                bakupWaitKeyDeley = bakupWaitKeyDeley + 1
            else:
                bakupWaitKeyDeley = maxSpeed
    elif key == 2424832: # 方向键 左， 后退
        fno = fno - 1
        cap.set(cv2.CAP_PROP_POS_FRAMES, fno)
    elif key == 2555904: # 方向键 右， 前进
        fno = fno + 1
        cap.set(cv2.CAP_PROP_POS_FRAMES, fno)
    elif key == ord('z'): # 撤回到上一个关键点
        fno = idx_and_labels.pop()[0]
        cap.set(cv2.CAP_PROP_POS_FRAMES, fno)
    elif key == ord('f'):
        fno = fno + 10
        cap.set(cv2.CAP_PROP_POS_FRAMES, fno)
    elif key == ord('d'):
        fno = fno - 10
        cap.set(cv2.CAP_PROP_POS_FRAMES, fno)
    elif key == 27: # 退出
        print('Current fno =',fno)
        cap.release()
        idx_and_labels.append([frame_num,-1])
        f = open(savePath,'w')
        for i in range(len(idx_and_labels) - 1):
            label = idx_and_labels[i][1]
            for j in range(idx_and_labels[i+1][0] - idx_and_labels[i][0]):
                f.write('%d %d\n'%(j + idx_and_labels[i][0], label))
        f.close()
        exit(0)

idx_and_labels.append([frame_num,-1])
f = open(savePath,'w')
for i in range(len(idx_and_labels) - 1):
    label = idx_and_labels[i][1]
    for j in range(idx_and_labels[i+1][0] - idx_and_labels[i][0]):
        f.write('%d %d\n'%(j + idx_and_labels[i][0], label))
f.close()